import os
import numpy as np

from tqdm import tqdm


def load_data(anno_dir):
    print('load data...')
    file_names = os.listdir(anno_dir)
    boxes = []
    for file_name in tqdm(file_names):
        file_path = os.path.join(anno_dir, file_name)
        with open(file_path) as f:
            lines = f.readlines()
            for line in lines:
                box = line.split()
                box = [float(box[3]), float(box[4])]
                boxes.append(box)
    return np.array(boxes)


def cas_iou(box, cluster):
    x = np.minimum(cluster[:, 0], box[0])
    y = np.minimum(cluster[:, 1], box[1])

    intersection = x * y
    area1 = box[0] * box[1]

    area2 = cluster[:, 0] * cluster[:, 1]
    iou = intersection / (area1 + area2 - intersection)

    return iou


def avg_iou(box, cluster):
    return np.mean([np.max(cas_iou(box[i], cluster)) for i in range(box.shape[0])])


def kmeans(box, k):
    # 取出一共有多少框
    row = box.shape[0]

    # 每个框各个点的位置
    distance = np.empty((row, k))

    # 最后的聚类位置
    last_clu = np.zeros((row,))

    np.random.seed()

    # 随机选5个当聚类中心
    cluster = box[np.random.choice(row, k, replace=False)]
    # cluster = random.sample(row, k)
    print('Clustering...')
    while True:
        # 计算每一行距离五个点的iou情况。
        for i in tqdm(range(row)):
            distance[i] = 1 - cas_iou(box[i], cluster)

        # 取出最小点
        near = np.argmin(distance, axis=1)

        if (last_clu == near).all():
            break

        # 求每一个类的中位点
        for j in range(k):
            cluster[j] = np.median(
                box[near == j], axis=0)

        last_clu = near

    return cluster


if __name__ == '__main__':
    SIZE = 800
    anchors_num = 9
    # 使用k聚类算法
    data = load_data(anno_dir=r'F:\dataset\object_detection\dota1.5\rotated_data\train_data\1024x1024\labels\train')
    out = kmeans(data, anchors_num)
    out = out[np.argsort(out[:, 0])]
    print('acc:{:.2f}%'.format(avg_iou(data, out) * 100))
    print(out * SIZE)
    data = (out * SIZE).astype(int)
    f = open("yolo_anchors.txt", 'w')
    row = np.shape(data)[0]
    print('save anchors ...')
    for i in tqdm(range(0, row, 3)):
        x_y = f'[{data[i][0]},{data[i][1]}, {data[i + 1][0]},{data[i + 1][1]}, {data[i + 2][0]},{data[i + 2][1]}]\n'
        f.write(x_y)
    f.close()
    print('complete!')
